//
// Created by cybaster on 25-6-7.
//

#include "ocs2_humanoid_robot_controller/controller/HumanoidRobotWBCController.h"

HumanoidRobotWBCController::HumanoidRobotWBCController(
    LeggedRobotInterface &leggedRobotInterface)
    : HumanoidRobotController(leggedRobotInterface) {}

void HumanoidRobotWBCController::setWBCController(
    const std::shared_ptr<legged::WbcBase> &wbcController) {
  wbc_ = wbcController;
}

vector_t HumanoidRobotWBCController::calculateJointAccAndTorque(
    vector_t optimizedState, vector_t optimizedInput, size_t plannerMode) {
  if (wbc_) {
    auto x = wbc_->update(optimizedState, optimizedInput, rbdState_,
                          plannerMode, timestep);
    vector_t jointAccAndTorque(info_.actuatedDofNum * 2);
    jointAccAndTorque.head(info_.actuatedDofNum) =
        x.segment(6, info_.actuatedDofNum);
    jointAccAndTorque.tail(info_.actuatedDofNum) = x.tail(info_.actuatedDofNum);
    return jointAccAndTorque;
  }

  return HumanoidRobotController::calculateJointAccAndTorque(
      optimizedState, optimizedInput, plannerMode);
}
